"""
数据增强
"""

import numpy as np
from paddle.vision.transforms import RandomCrop, BrightnessTransform
from PIL import Image
import cv2


def random_crop(img, crop_size, labels=None) -> np.ndarray:
    assert img.shape[0] >= crop_size[0] and img.shape[1] >= crop_size[1]
    st_i = np.random.randint(img.shape[0] - crop_size[0] + 1)
    st_j = np.random.randint(img.shape[1] - crop_size[1] + 1)
    transformed_labels = []
    if labels != None:
        for label in labels:
            xmin = np.clip(label[0] - st_i, 0, crop_size[0])
            ymin = np.clip(label[1] - st_j, 0, crop_size[1])
            xmax = np.clip(label[2] - st_i, 0, crop_size[0])
            ymax = np.clip(label[3] - st_j, 0, crop_size[1])
            transformed_labels.append([xmin, ymin, xmax, ymax])
    return img[st_i: st_i + crop_size[0], st_j: st_j + crop_size[1], :], np.array(transformed_labels).astype(int)


def random_transform(img, crop_size=(320, 640, 3), brightness=0.4, labels=None) -> np.ndarray:
    """
    实现图像的裁剪、亮度增强。

    Args:
    ---
        - img: 要变换的图片。
        - crop_size: 裁剪大小。
        - brightness: 亮度增强的参数。
    """
    # 图像裁剪
    img, labels = random_crop(img, crop_size[:2], labels=labels)
    # 亮度增强
    brightness_transform = BrightnessTransform(brightness)
    img = brightness_transform(img)

    return np.array(img), labels


# 测试变换功能
if __name__ == '__main__':
    import os
    imgs = os.listdir('./data/training_images/')
    img_file = None
    for i in imgs:
        if len(i) > 4 and i[-4:] == '.jpg':
            img_file = i
            break
    if img_file is None:
        raise ValueError('No images in target folder.')
    img = cv2.imread(os.path.join('./data/training_images/', img_file))
    img = random_transform(img)
    cv2.imshow('After Transform', img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()
